import os
import requests
import traceback
import time

class BaiduAIService:
    def __init__(self):
        self.API_KEY = os.getenv('BAIDU_API_KEY')
        self.SECRET_KEY = os.getenv('BAIDU_SECRET_KEY')
        self.access_token = self._get_access_token()

    def _get_access_token(self):
        """获取百度 AI 访问令牌"""
        url = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={self.API_KEY}&client_secret={self.SECRET_KEY}"
        response = requests.get(url)
        return response.json().get('access_token')

    def generate_emoji(self, text, style):
        """
        使用百度 AI 作画-高级版 API 生成表情包
        :param text: 生成提示词
        :param style: 生成风格
        :return: 生成的图片数据
        """
        try:
            print(f"\n=== Baidu AI Request ===")
            print(f"Text: {text}")
            print(f"Style: {style}")

            # 根据不同风格构建提示词
            style_prompts = {
                'cartoon': '卡通风格，可爱，表情包',
                'realistic': '写实风格，细节丰富，高质量',
                'artistic': '艺术风格，创意，充满活力',
                'anime': '动漫风格，日系，可爱',
                'comic': '漫画风格，粗线条，表现力强'
            }

            # 构建完整提示词
            prompt = f"{text}，{style_prompts.get(style, style_prompts['cartoon'])}"

            # 第一步：提交生成请求
            submit_url = f"https://aip.baidubce.com/rpc/2.0/ernievilg/v1/txt2imgv2?access_token={self.access_token}"
            
            payload = {
                "prompt": prompt,
                "width": 1024,
                "height": 1024,
                "style": style,
                "num": 1,
                "image_num": 1
            }

            headers = {
                'Content-Type': 'application/json',
                'Accept': 'application/json'
            }

            # 发送生成请求
            submit_response = requests.post(submit_url, headers=headers, json=payload)
            submit_result = submit_response.json()
            
            print("Submit response:", submit_result)

            # 修改这里：检查 task_id 的位置
            task_id = submit_result.get('data', {}).get('task_id')
            if task_id:
                # 第二步：轮询获取结果
                query_url = f"https://aip.baidubce.com/rpc/2.0/ernievilg/v1/getImgv2?access_token={self.access_token}"
                query_payload = {"task_id": task_id}
                
                # 最多等待 60 秒
                for _ in range(20):  # 增加重试次数
                    query_response = requests.post(query_url, headers=headers, json=query_payload)
                    query_result = query_response.json()
                    
                    print("Query response:", query_result)
                    
                    if 'data' in query_result:
                        final_image_list = query_result['data'].get('sub_task_result_list', [])[0].get('final_image_list', [])
                        if final_image_list:
                            image_url = final_image_list[0].get('img_url')
                            if image_url:
                                return {
                                    'success': True,
                                    'image_url': image_url
                                }
                    # 等待时间改为 3 秒
                    time.sleep(3)
                
                return {
                    'success': False,
                    'error_msg': '生成超时，请重试'
                }
            else:
                error_msg = submit_result.get('error_msg', '未知错误')
                print(f"Error: {error_msg}")
                return {
                    'success': False,
                    'error_msg': error_msg
                }

        except Exception as e:
            print(f"Baidu AI Service error: {str(e)}")
            print(traceback.format_exc())
            return {
                'success': False,
                'error_msg': str(e)
            }